#!/usr/bin/env python3
import os
import argparse
import gzip
import pysam
from bx.intervals.intersection import IntervalTree


class Interval:
    def __init__(self, chrom=None, start=None, end=None, name=None, strand=None):
        self.chrom = chrom
        self.chromStart = start
        self.chromEnd = end
        self.name = name
        self.strand = strand
        self.info = dict()

    __slots__ = ("chrom", "chromStart", "chromEnd", "name", "strand", "info")

    def __hash__(self):
        return hash((self.chrom, self.chromStart, self.chromEnd, self.strand))

    def __eq__(self, other):
        if self.chrom != other.chrom:
            return False
        if self.chromStart != other.chromStart:
            return False
        if self.chromEnd != other.chromEnd:
            return False
        if self.strand != other.strand:
            return False
        return True

    def __str__(self):
        return "{chrom}:{start}-{end} ({strand})".format(
            chrom=self.chrom, start=self.chromStart, end=self.chromEnd, strand=self.strand
        )

    __repr__ = __str__

    def set_info(self, **kwargs):
        for key, value in kwargs.items():
            self.info[key] = value

    def get_info(self, key):
        return self.info[key]

    def is_overlap(self, other):
        if self.chrom != other.chrom:
            return False
        if self.strand != other.strand:
            return False
        if self.chromStart > other.chromEnd:
            return False
        if other.chromStart > self.chromEnd:
            return False
        return True


class IvalTree:
    def __init__(self, strand=None):
        self.strand = strand
        self.__trees = dict()

    def add(self, obj, chrom, start, end, strand=None):
        if self.strand:
            key = (chrom, strand)
        else:
            key = chrom
        self.__add_record(key, start, end, obj)

    def __add_record(self, key, start, end, record):
        if key not in self.__trees.keys():
            self.__trees[key] = IntervalTree()
        self.__trees[key].insert(start, end, record)

    def add_record(self, record):
        assert isinstance(record, Interval)
        if self.strand:
            key = (record.chrom, record.strand)
        else:
            key = record.chrom
        start = record.chromStart
        end = record.chromEnd
        self.__add_record(key, start, end, record)

    def find(self, chrom, start, end, strand=None):
        if self.strand:
            key = (chrom, strand)
        else:
            key = chrom
        overlaped_rec = list()
        if key in self.__trees.keys():
            overlaped_rec = self.__trees[key].find(start, end)
        return overlaped_rec

    def iter_ival(self, chrom, strand=None):
        if self.strand:
            key = (chrom, strand)
        else:
            key = chrom
        nodes = list()
        tree = self.__trees[key]
        fn = nodes.append
        tree.traverse(fn)
        for node in nodes:
            yield node.interval

    def iter_all(self):
        for key, tree in sorted(self.__trees.items()):
            nodes = list()
            fn = nodes.append
            tree.traverse(fn)
            for node in nodes:
                yield node.interval


class NoChimeric:
    def __init__(self, read):
        self.read = read
        self.name = read.query_name.replace(":","_")
        self.sequence = read.get_forward_sequence()
        try:
            self.quality = pysam.qualities_to_qualitystring(
                read.get_forward_qualities())
        except TypeError:
            self.quality = 'F' * len(self.sequence)
        if not self.quality:
            self.quality = 'F' * len(self.sequence)
        self.strand = '-' if read.is_reverse else '+'
        self.cigar = read.cigartuples
        self.chrom = read.reference_name
        self.chromStart = read.reference_start

    def is_NoChimeric(self, knownJS):
        if self.read.is_qcfail:
            return False
        if self.read.reference_name is None:
            return False
        if self.read.reference_start is None:
            return False
        if self.read.reference_end is None:
            return False

        for cigar_tuple in self.cigar:
            if cigar_tuple[0] == 4 and cigar_tuple[1] >= 20:
                return False

        block = self.read.get_blocks()
        block_len = len(block)
        if block_len > 1:
            JS = []
            for indx in range(1, block_len):
                JS.append((block[indx-1][1], block[indx][0]))
            new_JS_li = []
            for js in JS:
                js_rec = Interval(self.chrom, js[0], js[1])
                overlap_js = knownJS.find(js_rec.chrom, js_rec.chromStart, js_rec.chromEnd)
                if not any([x == js_rec for x in overlap_js]):
                    new_JS_li.append(js_rec)
            if not new_JS_li:
                return True
            else:
                return False
        else:
            return True


def main(bamfile, known_JS, fastq):
    bam = pysam.AlignmentFile(bamfile, "rb")
    fq = gzip.open(fastq, "wb")

    js_tree = IvalTree()
    with open(known_JS, "r") as f:
        for line in f:
            chrom, start, end, name, score = line.split()[:5]
            rec = Interval(chrom=chrom, start=int(start), end=int(end))
            js_tree.add_record(rec)

    reads_count = no_js = 0
    for read in bam:
        if read.is_secondary:
            continue
        reads_count += 1
        align = NoChimeric(read)
        if not align.is_NoChimeric(js_tree):
            fq.write("@{}\n{}\n+\n{}\n".format(align.name, align.sequence, align.quality).encode())
        else:
            no_js += 1
    print("total: {} reads. no junction reads mapped: {}. rate: {}.".format(reads_count, no_js, no_js/reads_count))


def run():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-i", dest="bam",
                        type=str, required=True, help="bamfile")
    parser.add_argument("-k", dest="knownjs",
                        type=str, required=True, help="knownjs, cat intron_bed | sort -k 1,1 -k 2,2n -k 3,3n -k 6,6 -u > knownjs")
    parser.add_argument("-o", dest="fastq", default=os.devnull,
                        type=str, help="output fastq")
    args = parser.parse_args()
    
    main(args.bam, args.knownjs, args.fastq)

    
if __name__ == "__main__":
    run()
